

import torch


def generate_uniform_probs(
    num_tokens: int,
    num_draft_tokens: list[int],
    max_spec_num: int,
) -> torch.Tensor:
    batch_size = len(num_draft_tokens)
    
    assert sum(num_draft_tokens) == num_tokens
    uniform_probs = torch.rand((num_tokens, ), dtype=torch.float32,)  # 使用torch替代np.random
    
    segments = torch.split(uniform_probs, num_draft_tokens)
    result = torch.full((batch_size, max_spec_num), -1.0, dtype=torch.float32)

    for i, seg in enumerate(segments):
        n = min(len(seg), max_spec_num)
        result[i, :n] = seg[:n]
    
    return result

x = [[123, 234, -1], [222, -1, -1]]
print(x[0]!=-1)
